//
// Created by fred on 2020/4/14.
//
#include <jni.h>
#include <string>
#include "Log.h"
#include "Key.h"
#include "MD5.h"
#include "Verification.h"
#include "base64.h"

extern "C"
JNIEXPORT jstring JNICALL
Java_com_pg_lockly_key_KeyManager_a(JNIEnv *env, jobject, jobject contextObj) {
    std::string endKey = "u6Z16aR1weZck7o5raI8M8w994Ce36Id";
    return checkAndReturn32(env, AES_KEY_START, endKey, contextObj);
}

extern "C"
JNIEXPORT jstring JNICALL
Java_com_pg_lockly_key_KeyManager_b(JNIEnv *env, jobject, jobject context) {
    std::string endKey = "ce7gFjUoKmDzp4D719RnWRwK3OpZCT3B";
    return checkAndReturn32(env, DES_KEY_START, endKey, context);
}


extern "C"
JNIEXPORT jstring JNICALL
Java_com_pg_lockly_key_KeyManager_c(JNIEnv *env, jobject, jobject context, jstring input) {
    if (!verification(env, context)) {
        throwException(env, "Unknown Error");
        return nullptr;
    }
    if (context == nullptr || input == nullptr){
        return nullptr;
    }
    int ENCRYPT_DATA_BLOCK_SIZE = 64;
    std::string endKey = "DJQ6E4BxAJQJU0A4Ao49GFNWA6DPCDBf";
    std::string sStr = getReleaseSign(RSA_PUBLIC_KEY_START, endKey) + RSA_PUBLIC_KEY_SURPLUS;
    const char *publicKeyStr = sStr.c_str();
    char dst[Base64decode_len(publicKeyStr)];
    int size = Base64decode(dst, publicKeyStr);
    jbyteArray key = env->NewByteArray(size);
    env->SetByteArrayRegion(key, 0, size, (const jbyte *) (dst));
    env->ReleaseByteArrayElements(key, env->GetByteArrayElements(key, JNI_FALSE), 0);
    jclass jStringClass = env->FindClass("java/lang/String");
    if (jStringClass == nullptr) {
        return input;
    }
    jmethodID jStringBytesId = env->GetMethodID(jStringClass, "getBytes", "()[B");
    env->DeleteLocalRef(jStringClass);
    if (jStringBytesId == nullptr) {
        return input;
    }
    auto in = (jbyteArray)(env->CallObjectMethod(input, jStringBytesId));
    jclass x509KeySpecClass = env->FindClass("java/security/spec/X509EncodedKeySpec");
    if (x509KeySpecClass == nullptr) {
        return input;
    }
    jmethodID initMethodId = env->GetMethodID(x509KeySpecClass, "<init>", "([B)V");
    if (initMethodId == nullptr) {
        return input;
    }
    jobject x509Key = env->NewObject(x509KeySpecClass, initMethodId, key);
    env->DeleteLocalRef(key);
    env->DeleteLocalRef(x509KeySpecClass);
    jclass factoryClass = env->FindClass("java/security/KeyFactory");
    if (factoryClass == nullptr) {
        return input;
    }
    jmethodID getInstanceId = env->GetStaticMethodID(factoryClass, "getInstance",
                                                     "(Ljava/lang/String;)Ljava/security/KeyFactory;");
    if (getInstanceId == nullptr) {
        return input;
    }
    const char *keyAlgorithm = "RSA";
    jstring KEY_ALGORITHM = env->NewStringUTF(keyAlgorithm);
    jobject keyFactory = env->CallStaticObjectMethod(factoryClass, getInstanceId, KEY_ALGORITHM);
    jmethodID generatePublicId = env->GetMethodID(factoryClass, "generatePublic",
                                                  "(Ljava/security/spec/KeySpec;)Ljava/security/PublicKey;");
    env->DeleteLocalRef(KEY_ALGORITHM);
    if (generatePublicId == nullptr) {
        return input;
    }
    jobject publicKey = env->CallObjectMethod(keyFactory, generatePublicId, x509Key);
    env->DeleteLocalRef(x509Key);
    jmethodID getAlgorithmId = env->GetMethodID(factoryClass, "getAlgorithm",
                                                "()Ljava/lang/String;");
    env->DeleteLocalRef(factoryClass);
    if (getAlgorithmId == nullptr) {
        return input;
    }
    auto algorithm = (jstring)(env->CallObjectMethod(keyFactory, getAlgorithmId));
    env->DeleteLocalRef(keyFactory);
    jclass bcpClass = env->FindClass("org/bouncycastle/jce/provider/BouncyCastleProvider");
    if (bcpClass == nullptr) {
        return input;
    }
    initMethodId = env->GetMethodID(bcpClass, "<init>", "()V");
    if (initMethodId == nullptr) {
        return input;
    }
    jobject bcp = env->NewObject(bcpClass, initMethodId);
    env->DeleteLocalRef(bcpClass);
    jclass cipherClass = env->FindClass("javax/crypto/Cipher");
    if (cipherClass == nullptr) {
        return input;
    }
    getInstanceId = env->GetStaticMethodID(cipherClass, "getInstance",
                                           "(Ljava/lang/String;Ljava/security/Provider;)Ljavax/crypto/Cipher;");
    if (getInstanceId == nullptr) {
        return input;
    }
    jobject cipher = env->CallStaticObjectMethod(cipherClass, getInstanceId, algorithm, bcp);
    env->DeleteLocalRef(algorithm);
    env->DeleteLocalRef(bcp);
    jmethodID initId = env->GetMethodID(cipherClass, "init", "(ILjava/security/Key;)V");
    if (initId == nullptr) {
        return input;
    }
    env->CallVoidMethod(cipher, initId, 1, publicKey);
    env->DeleteLocalRef(publicKey);
    jmethodID doFinalId = env->GetMethodID(cipherClass, "doFinal", "([BII)[B");
    env->DeleteLocalRef(cipherClass);
    if (doFinalId == nullptr) {
        return input;
    }
    int inLen = env->GetArrayLength(in);
    int offset = 0, i = 0;
    jclass outputClass = env->FindClass("java/io/ByteArrayOutputStream");
    if (outputClass == nullptr) {
        return input;
    }
    initMethodId = env->GetMethodID(outputClass, "<init>", "()V");
    if (initMethodId == nullptr) {
        return input;
    }
    jobject out = env->NewObject(outputClass, initMethodId);
    jmethodID outWriteId = env->GetMethodID(outputClass, "write", "([BII)V");
    if (outWriteId == nullptr) {
        return input;
    }
    jbyteArray cache = nullptr;
    while (inLen - offset > 0) {
        if (inLen - offset > ENCRYPT_DATA_BLOCK_SIZE) {
            cache = (jbyteArray)(env->CallObjectMethod(cipher, doFinalId, in, offset,
                                                                  ENCRYPT_DATA_BLOCK_SIZE));
        } else {
            cache = (jbyteArray)(env->CallObjectMethod(cipher, doFinalId, in, offset,
                                                                  inLen - offset));
        }
        int cacheLen = env->GetArrayLength(cache);
        env->CallVoidMethod(out, outWriteId, cache, 0, cacheLen);
        i++;
        offset = i * ENCRYPT_DATA_BLOCK_SIZE;
    }
    env->DeleteLocalRef(cipher);
    env->DeleteLocalRef(cache);
    jmethodID outToByteArray = env->GetMethodID(outputClass, "toByteArray", "()[B");
    if (outToByteArray == nullptr) {
        return input;
    }
    jmethodID outClose = env->GetMethodID(outputClass, "close", "()V");
    env->DeleteLocalRef(outputClass);
    if (outClose == nullptr) {
        return input;
    }
    auto result = (jbyteArray)(env->CallObjectMethod(out, outToByteArray));
    env->CallVoidMethod(out, outClose);
    env->DeleteLocalRef(out);
    int srcSize = env->GetArrayLength(result);
    int base64Len = Base64encode_len(srcSize);
    char outBase64[base64Len + 1];
    const char *src = reinterpret_cast<const char *>(env->GetByteArrayElements(result, JNI_FALSE));
    env->DeleteLocalRef(result);
    Base64encode(outBase64, src, srcSize);
    outBase64[base64Len] = '\0';
    //input is not valid Modified UTF-8: illegal start byte 0x81
    return env->NewStringUTF(outBase64);
}

extern "C"
JNIEXPORT jstring JNICALL
Java_com_pg_lockly_key_KeyManager_d(JNIEnv *env, jobject , jobject context, jstring input) {
    if (!verification(env, context)) {
        throwException(env, "Unknown Error");
        return nullptr;
    }
    if (context == nullptr || input == nullptr){
        return nullptr;
    }
    const char *d = env->GetStringUTFChars(input, JNI_FALSE);
    int inLen = Base64decode_len(d);
    char plain[inLen];
    int realLen = Base64decode(plain, d);
    jbyteArray in = env->NewByteArray(realLen);
    env->SetByteArrayRegion(in, 0, realLen, (const jbyte *) (plain));
    env->ReleaseByteArrayElements(in, env->GetByteArrayElements(in, JNI_FALSE), 0);
    int CRYPT_DATA_BLOCK_SIZE = 128;
    std::string endKey = "h1kFihGP92wx0FBlAkQpE+FyAyA+ShCs";
    std::string sStr = getReleaseSign(RSA_PRIVATE_KEY_START, endKey) + RSA_PRIVATE_KEY_SURPLUS;
    const char *privateKeyStr = sStr.c_str();
    int pStrLen = Base64decode_len(privateKeyStr);
    char dst[pStrLen];
    int size = Base64decode(dst, privateKeyStr);
    jbyteArray key = env->NewByteArray(size);
    env->SetByteArrayRegion(key, 0, size, (const jbyte *) (dst));
    env->ReleaseByteArrayElements(key, env->GetByteArrayElements(key, JNI_FALSE), 0);
    jclass pkcsClass = env->FindClass("java/security/spec/PKCS8EncodedKeySpec");
    if (pkcsClass == nullptr) {
        return input;
    }
    jmethodID initMethodId = env->GetMethodID(pkcsClass, "<init>", "([B)V");
    if (initMethodId == nullptr) {
        return input;
    }
    jobject pkcsSpec = env->NewObject(pkcsClass, initMethodId, key);
    env->DeleteLocalRef(pkcsClass);
    jclass factoryClass = env->FindClass("java/security/KeyFactory");
    if (factoryClass == nullptr) {
        return input;
    }
    jmethodID getInstanceId = env->GetStaticMethodID(factoryClass, "getInstance",
                                                     "(Ljava/lang/String;)Ljava/security/KeyFactory;");
    if (getInstanceId == nullptr) {
        return input;
    }
    const char *keyAlgorithm = "RSA";
    jstring KEY_ALGORITHM = env->NewStringUTF(keyAlgorithm);
    jobject keyFactory = env->CallStaticObjectMethod(factoryClass, getInstanceId, KEY_ALGORITHM);
    jmethodID generatePrivateId = env->GetMethodID(factoryClass, "generatePrivate",
                                                   "(Ljava/security/spec/KeySpec;)Ljava/security/PrivateKey;");
    env->DeleteLocalRef(KEY_ALGORITHM);
    if (generatePrivateId == nullptr) {
        return input;
    }
    jobject privateKey = env->CallObjectMethod(keyFactory, generatePrivateId, pkcsSpec);
    env->DeleteLocalRef(pkcsSpec);
    jmethodID getAlgorithmId = env->GetMethodID(factoryClass, "getAlgorithm",
                                                "()Ljava/lang/String;");
    env->DeleteLocalRef(factoryClass);
    if (getAlgorithmId == nullptr) {
        return input;
    }
    auto algorithm = (jstring)(env->CallObjectMethod(keyFactory, getAlgorithmId));
    env->DeleteLocalRef(keyFactory);
    jclass bcpClass = env->FindClass("org/bouncycastle/jce/provider/BouncyCastleProvider");
    if (bcpClass == nullptr) {
        return input;
    }
    initMethodId = env->GetMethodID(bcpClass, "<init>", "()V");
    if (initMethodId == nullptr) {
        return input;
    }
    jobject bcp = env->NewObject(bcpClass, initMethodId);
    env->DeleteLocalRef(bcpClass);
    jclass cipherClass = env->FindClass("javax/crypto/Cipher");
    if (cipherClass == nullptr) {
        return input;
    }
    getInstanceId = env->GetStaticMethodID(cipherClass, "getInstance",
                                           "(Ljava/lang/String;Ljava/security/Provider;)Ljavax/crypto/Cipher;");
    if (getInstanceId == nullptr) {
        return input;
    }
    jobject cipher = env->CallStaticObjectMethod(cipherClass, getInstanceId, algorithm, bcp);
    env->DeleteLocalRef(bcp);
    env->DeleteLocalRef(algorithm);
    jmethodID initId = env->GetMethodID(cipherClass, "init", "(ILjava/security/Key;)V");
    if (initId == nullptr) {
        return input;
    }
    env->CallVoidMethod(cipher, initId, 2, privateKey);
    env->DeleteLocalRef(privateKey);
    jmethodID doFinalId = env->GetMethodID(cipherClass, "doFinal", "([BII)[B");
    env->DeleteLocalRef(cipherClass);
    if (doFinalId == nullptr) {
        return input;
    }
    inLen = env->GetArrayLength(in);
    int offset = 0, i = 0;
    jclass outputClass = env->FindClass("java/io/ByteArrayOutputStream");
    if (outputClass == nullptr) {
        return input;
    }
    initMethodId = env->GetMethodID(outputClass, "<init>", "()V");
    if (initMethodId == nullptr) {
        return input;
    }
    jobject out = env->NewObject(outputClass, initMethodId);
    jmethodID outWriteId = env->GetMethodID(outputClass, "write", "([BII)V");
    jbyteArray cache = nullptr;
    while (inLen - offset > 0) {
        if (inLen - offset > CRYPT_DATA_BLOCK_SIZE) {
            cache = (jbyteArray)(env->CallObjectMethod(cipher, doFinalId, in, offset,
                                                                  CRYPT_DATA_BLOCK_SIZE));
        } else {
            cache = (jbyteArray)(env->CallObjectMethod(cipher, doFinalId, in, offset,
                                                                  inLen - offset));
        }
        int cacheLen = env->GetArrayLength(cache);
        env->CallVoidMethod(out, outWriteId, cache, 0, cacheLen);
        i++;
        offset = i * CRYPT_DATA_BLOCK_SIZE;
    }

    env->DeleteLocalRef(cache);
    env->DeleteLocalRef(cipher);
    jmethodID outToByteArray = env->GetMethodID(outputClass, "toByteArray", "()[B");
    if (outToByteArray == nullptr) {
        return input;
    }
    jmethodID outClose = env->GetMethodID(outputClass, "close", "()V");
    env->DeleteLocalRef(outputClass);
    if (outClose == nullptr) {
        return input;
    }
    auto result = (jbyteArray)(env->CallObjectMethod(out, outToByteArray));
    env->CallVoidMethod(out, outClose);
    env->DeleteLocalRef(out);
    jclass jStringClass = env->FindClass("java/lang/String");
    if (jStringClass == nullptr) {
        return input;
    }
    initId = env->GetMethodID(jStringClass, "<init>", "([B)V");
    if (initId == nullptr) {
        return input;
    }
    return (jstring)(env->NewObject(jStringClass, initId, result));
}

extern "C"
JNIEXPORT jstring JNICALL
Java_com_pg_lockly_key_KeyManager_e(JNIEnv *env, jobject , jobject context, jstring de) {
    if (context == nullptr || de == nullptr){
        return nullptr;
    }
    auto *source = (unsigned char *) env->GetStringUTFChars(de, JNI_FALSE);
    int length = strlen((char *) (source));
    AES_KEY key;
    int result = setAesKey(env, &key, AES_ENCRYPT, context);
    if (result < 0 || length < 1) {
        return nullptr;
    }
    unsigned char in[length + AES_BLOCK_SIZE];
    int groupNum = length / AES_BLOCK_SIZE + 1;
    int padding = AES_BLOCK_SIZE - length % AES_BLOCK_SIZE;
    int inSize = groupNum * AES_BLOCK_SIZE;
    memset((void *) in, padding, inSize);
    memcpy((void *) in, source, length);
    unsigned char out[inSize];
    Log("groupNum:%d realSize:%d padding:%d", groupNum, inSize, padding);
    for (int i = 0; i < groupNum; ++i) {
        AES_ecb_encrypt(in + i * AES_BLOCK_SIZE, out + i * AES_BLOCK_SIZE, &key,
                        AES_ENCRYPT);
    }

    int base64Len = Base64encode_len(inSize);
    char outBase64[base64Len + 1];
    int ber = Base64encode(outBase64, reinterpret_cast<const char *>(out), inSize);
    outBase64[base64Len] = '\0';
    Log("加密Len:%d  base64Len:%d  ber:%d result:%s", inSize, base64Len, ber, outBase64);
    return env->NewStringUTF(outBase64);
}

extern "C"
JNIEXPORT jstring JNICALL
Java_com_pg_lockly_key_KeyManager_f(JNIEnv *env, jobject, jobject context, jstring de) {
    if (context == nullptr || de == nullptr){
        return nullptr;
    }
    const char *temp = (const char *) env->GetStringUTFChars(de, JNI_FALSE);
    int base64Len = Base64decode_len(temp);
    char plain[base64Len];
    int length = Base64decode(plain, temp);
    auto *source = reinterpret_cast<unsigned char *>(plain);
    AES_KEY key;
    int result = setAesKey(env, &key, AES_DECRYPT, context);
    if (result < 0 || length < 1) {
        return nullptr;
    }
    if (length % AES_BLOCK_SIZE > 0) {
        throwException(env, "data not multiple of block length");
        return nullptr;
    }

    int groupNum = length / AES_BLOCK_SIZE;
    unsigned char out[length];
    for (int i = 0; i < groupNum; ++i) {
        AES_ecb_encrypt(source + i * AES_BLOCK_SIZE, out + i * AES_BLOCK_SIZE, &key,
                        AES_DECRYPT);
    }
    int paddingNum = out[length - 1];
    int realLen = length - paddingNum;
    char resultOut[realLen + 1];
    memcpy((void *) resultOut, out, realLen);
    resultOut[realLen] = '\0';
    return env->NewStringUTF(resultOut);
}


int setAesKey(JNIEnv *env, AES_KEY *st, const int enc, jobject context) {
    if (!verification(env, context)) {
        throwException(env, "Unknown Error");
        return -1;
    }
    std::string endKey = "Q2e/gFjdf8yoKEtp4D7M8X0Rw3QpZT3B";
    std::string temp = getReleaseSign(AES_DID_KEY_START, endKey);
    const char *aesUserKey = temp.c_str();
    if (!aesUserKey || !st) {
        throwException(env, "Key is null");
        return -1;
    }
//    Log("key is:%s", aesUserKey);
    int keyLen = temp.length();
//    int keyLen1 = strlen(aesUserKey);
//    Log("keyLen:%d   %d", keyLen,keyLen1);
    if (keyLen != 16 && keyLen != 24 && keyLen != 32) {
        throwException(env, "key size must be 16/24/32 bytes");
        return -3;
    }
    int bits = 128;
    if (keyLen == 24) {
        bits = 192;
    } else if (keyLen == 32) {
        bits = 256;
    }
    //CFB模式 解密也设置 AES_set_encrypt_key
    if (enc == AES_ENCRYPT) {
        return AES_set_encrypt_key(reinterpret_cast<const unsigned char *>(aesUserKey), bits, st);
    } else {
        return AES_set_decrypt_key(reinterpret_cast<const unsigned char *>(aesUserKey), bits, st);
    }

}

/**
 * 调用该方法之后就不能再使用 env 否则会有异常
 */
void throwException(JNIEnv *env, const char *errorInfo) {
    jclass jc = env->FindClass("com/pg/lockly/key/JniException");
    jboolean flag = env->ExceptionCheck();
    if (flag) {
        env->ExceptionDescribe();
        env->ExceptionClear();
    }
    env->ThrowNew(jc, errorInfo);
    env->DeleteLocalRef(jc);
}
